#pragma once

#include <string>
#include <mutex>
#include <atomic>
#include <cstdlib>
#include <stdexcept>
#include <memory>
#include <iostream>
#include <algorithm>
#include "util/allocator.hh"
#include "singleton.hpp"

#define USE_POOL 1

namespace kratos { namespace corelib {

class _ObjectPoolBase {
public:
    virtual ~_ObjectPoolBase() {}
    // Returns the allocated object's count
    // @return The allocated object's count
    virtual std::size_t getCount() const = 0;
    // Returns std::type_info of 'ObjectType'
    // @return std::type_info of 'ObjectType'
    virtual const std::type_info& getTypeInfo() const = 0;
    // Returns the size in byte of 'ObjectType'
    virtual std::size_t getObjectSize() const = 0;
};

using PoolList = std::list<_ObjectPoolBase*>;

// Object pool manager
class _ObjectPoolManager {    
    PoolList poolList_; // object pool list
    std::mutex lock_; // thread lock

public:
    ~_ObjectPoolManager() {
        PoolList temp = poolList_;
        for (auto pool : temp) {
            delete pool;
        }
    }
    // Returns pools
    const PoolList getPools() {
        std::lock_guard<std::mutex> lock(lock_);
        return poolList_;
    }

private:
    // Adds a pool singleton pointer
    void add(_ObjectPoolBase* pool) {
        std::lock_guard<std::mutex> lock(lock_);
        poolList_.push_back(pool);
    }

    void remove(_ObjectPoolBase* pool) {
        std::lock_guard<std::mutex> lock(lock_);
        poolList_.erase(std::remove(poolList_.begin(), poolList_.end(), pool));
    }

private:
    template <typename ObjectType, size_t Count>
    friend class _ObjectPool;
};

#define ObjectPoolManagerRef (*kratos::corelib::Singleton<kratos::corelib::_ObjectPoolManager>::instance())
#define ObjectPoolManagerSinglton kratos::corelib::Singleton<kratos::corelib::_ObjectPoolManager>

// object pool
// ObjectType object type
// Count pre-allocated count of object for each expand
template <typename ObjectType, size_t Count = 64>
class _ObjectPool : public _ObjectPoolBase {
    _ObjectPool(const _ObjectPool&) = delete;
    _ObjectPool(_ObjectPool&&) = delete;
    _ObjectPool<ObjectType>& operator=(const _ObjectPool&) = delete;

private:
    using FreeObjects = std::vector<ObjectType*>;
    using MemoryTrunks = std::list<char*>;
    FreeObjects freeObjects_; // free object list
    std::size_t freeCount_; // The count of object in free list
    MemoryTrunks memoryTrunks_; // occupied memory block list
    std::atomic<size_t> allocatedCount_; // occupied object count
    std::atomic<size_t> objectInHeap_; // The object count in heap
    std::mutex lock_; // thread lock

public:
    // ctor
    _ObjectPool() : freeCount_(0), allocatedCount_(0), objectInHeap_(0) {
        std::lock_guard<std::mutex> lock(lock_);
        _newTrunk();
        ObjectPoolManagerRef.add(this);
    }

    // dtor
    virtual ~_ObjectPool() {
        for (auto trunk : memoryTrunks_) {
            free(trunk);
        }
        ObjectPoolManagerRef.remove(this);
    }

    // Create a instance of 'ObjectType'
    // @param args The arguments for ctor of ObjectType
    // @return The instance pointer of ObjectType
    template <typename...Args>
    ObjectType* createObject(Args...args) {
#if !USE_POOL
        allocatedCount_ += 1;
        objectInHeap_ += 1;
        return new ObjectType(args...);
#else
        ObjectType* memoryPtr = nullptr;
        allocatedCount_.fetch_add(1);
        {
            std::lock_guard<std::mutex> lock(lock_);
            if (!freeCount_) {                
                if (!_newTrunk()) {
                    return nullptr;
                }
            }
            auto index = freeCount_ - 1;
            memoryPtr = freeObjects_[index];
            freeCount_ = index;
        }
        auto objectPtr = new(memoryPtr)ObjectType(args...);
        return objectPtr;
#endif // !USE_POOL
    }

    // Recycling
    // @param objectPtr The pointer of 'ObjectType'
    void disposeObject(ObjectType* objectPtr) {
#if !USE_POOL
        objectInHeap_ -= 1;
        allocatedCount_ -= 1;
        delete objectPtr;
#else
        allocatedCount_.fetch_sub(1);
        objectPtr->~ObjectType();
        std::lock_guard<std::mutex> lock(lock_);
        if (freeCount_ >= freeObjects_.size()) {
            throw std::runtime_error("Try to free object in pool but never allocated!");
            return;
        }
        freeObjects_[freeCount_] = objectPtr;
        freeCount_ += 1;
#endif // !USE_POOL
    }

    // Returns the allocated object's count
    // @return The allocated object's count
    virtual size_t getCount() const override {
        return objectInHeap_;
    }

    // Returns std::type_info of 'ObjectType'
    // @return std::type_info of 'ObjectType'
    virtual const std::type_info& getTypeInfo() const override {
        return typeid(ObjectType);
    }

    // Returns the size in byte of 'ObjectType'
    virtual std::size_t getObjectSize() const override {
        return sizeof(ObjectType);
    }

private:
    // Pre-allocates more objects
    bool _newTrunk() {
        auto memory = (char*)malloc(Count * sizeof(ObjectType));
        if (!memory) {
            return false;
        }
        if (!freeObjects_.size()) {
            freeObjects_.resize(Count);
        } else {
            freeObjects_.resize(freeObjects_.size() + Count);
        }
        for (auto i = 0; i < Count; i++) {
            freeObjects_[i] = reinterpret_cast<ObjectType*>(memory + i * sizeof(ObjectType));
        }
        objectInHeap_.fetch_add(Count);
        memoryTrunks_.push_back(memory);
        freeCount_ += Count;
        return true;
    }
};

template <typename ObjectType>
static inline void DestroyObjectPool() {
    kratos::corelib::Singleton<kratos::corelib::_ObjectPool<ObjectType>>::destroy();
}

template <typename ObjectType>
static inline kratos::corelib::_ObjectPool<ObjectType>* getObjectPoolPtr() {
    return kratos::corelib::Singleton<kratos::corelib::_ObjectPool<ObjectType>>::instance();
}

template <typename ObjectType>
static inline kratos::corelib::_ObjectPool<ObjectType>& getObjectPoolRef() {
    return *(kratos::corelib::Singleton<kratos::corelib::_ObjectPool<ObjectType>>::instance());
}

// Returns a object from pool and constructs it
// ObjectType The prototype of object
// Args The argument of object's constructor
// @return The new created object's pointer
template <typename ObjectType, typename...Args>
static inline ObjectType* allocate(Args...args) {
    return getObjectPoolRef<ObjectType>().createObject(args...);
}

// Returns a object from pool and constructs it
// ObjectType The prototype of object
// Args The argument of object's constructor
// @return The new created object's reference
template <typename ObjectType, typename...Args>
static inline ObjectType& allocateRef(Args...args) {
    auto object = getObjectPoolRef<ObjectType>().createObject(args...);
    if (!object) {
        throw std::runtime_error("out of memory");
    }
    return *object;
}

// Recycles to object pool
// @param objectPtr Object pointer
template <typename ObjectType>
static inline void dispose(ObjectType* objectPtr) {
    getObjectPoolRef<ObjectType>().disposeObject(objectPtr);
}

// Recycles to object pool
// @param objectPtr Object reference
template <typename ObjectType>
static inline void dispose(ObjectType& objectRef) {
    dispose(&objectRef);
}

// Destructor for smart pointer
template <typename T>
class Deleter {
public:
    void operator()(T* obj) {
        dispose(obj);
    }
};

template <typename T> using unique_pool_ptr = std::unique_ptr<T, Deleter<T>>;

}}

using ObjectPoolManager = kratos::corelib::_ObjectPoolManager;

#define ObjectPoolSingleton(T) kratos::corelib::Singleton<kratos::corelib::_ObjectPool<T>>
